% clear;clc;
% %% Preliminary
% load('Result.mat');
% load('IRF.mat');
% Folder = ['TableGraphs\\IRF_Decomposition\\'];

function Step_2_Decomposition(SOLUTION,Folder)
% load([File,'.mat']);
mkdir(Folder);

PP          =   SOLUTION.PP;
SS          =   SOLUTION.SS;

MODEL       =   SOLUTION.MODEL;
solution    =   SOLUTION.gxhx;
EquJac      =   SOLUTION.EquJac;
%% Generate the IRF
T           =   100;
StdVec      =   [1 1 1 -1 -1 60 -5]'*0.01/4;
IRF         =   IRF_1order(MODEL,solution,StdVec,[],T);
%% Decompose the IRF
DecDev      =   TempFun_IrfDec(PP,SS,MODEL,EquJac,IRF);

%% Group-wise Decomposition
ShockList   =   fieldnames(DecDev);
% Setup
VarGroup    =   struct('Interest_Dom',{{'ir'}},'Interest_Ext',{{'ir_ext'}}, ...
                       'w_H',{{'w_H'}},'w_N',{{'w_N'}},...
                       'Inflation',{{'Pir'}},'ExchangeRate',{{'dE'}},...
                       'T',{{'T'}},'Profit',{{'Profit'}}, ...
                       'TAU',{{'TAU'}},'Distribution',{{'Dist'}});
VarGroupLabel=  struct('Interest_Dom','Interest Rate, Domestic',...
                       'Interest_Ext','Interest Rate, Foreign', ...
                       'w_H','Wage, Home',...
                       'w_N','Wage, Non-tradable', ...
                       'Inflation','Inflation', ...
                       'ExchangeRate','Exchange Rate',...
                       'T','Transfer',...
                       'Profit','Dividend',...
                       'TAU','Tax',...
                       'Distribution','Distribution'...
                       );
GroupList   =   fields(VarGroup);
GroupLabelList= cell(length(GroupList),1);
for ii=1:length(GroupList)
    GroupLabelList{ii}  =   VarGroupLabel.(GroupList{ii});
end
GroupDecDev =   struct();
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    TempList    =   fields(DecDev.(SH));
    for jj=1:length(TempList)
        VV          =   TempList{jj};
        GroupDecDev.(SH).(VV) ...
                    =   zeros(MODEL.VAR.(VV).Dim,T,length(GroupList));
        for gg=1:length(GroupList)
            GG          =   GroupList{gg};
            GGTempList  =   VarGroup.(GG);
            for kk=1:length(GGTempList)
                GroupDecDev.(SH).(VV)(:,:,gg)    =   GroupDecDev.(SH).(VV)(:,:,gg)+DecDev.(SH).(VV).(GGTempList{kk});
            end
        end
    end
end

%% Collect the Results
TempVarList =   {'C','AggStat'};
IrfAgg      =   struct();
IrfDec      =   struct();
IrfDecTable =   struct();
Table_Rows  =   GroupList;
Table_Cols  =   {'C';...
                 'dom';'ext';'N';'H';'poor';'rich';...
                 'dom_N';'dom_H';'ext_N';'ext_H';...
                 'dom_poor';'dom_rich';'ext_poor';'ext_rich';...
                 'H_poor';'H_rich';'N_poor';'N_rich'};
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    % Peak IRF of Aggregate C
    [MaxVal,MaxIdx] = TempFun_MaxC(IRF.(SH).C);
    AggInd      =   MaxIdx;
    Normalizer  =   MaxVal;
    % [~,AggInd]  =   max(abs(IRF.(SH).C));
    % First 8 Quarters
%     AggInd      =   (1:8);
    % Normalizer
%     Normalizer  =   abs(sum(IRF.(SH).C(:,AggInd)));
    for jj=1:length(TempVarList)
        vv          =   TempVarList{jj};
        TempSize    =   size(GroupDecDev.(SH).(vv));
        IrfAgg.(SH).(vv) ...
                    =   reshape(sum(GroupDecDev.(SH).(vv)(:,AggInd,:),2),TempSize([1,3])) ...
                        /Normalizer;
        IrfDec.(SH).(vv) ...
                    =   IrfAgg.(SH).(vv)./sum(IrfAgg.(SH).(vv),2);
    end
    %----------------------------------------------------------------------
    % Tables
    %----------------------------------------------------------------------
    % Aggregated
    Table_Data  =   [IrfAgg.(SH).C;IrfAgg.(SH).AggStat]';
    IrfAggTable.(SH) ...
                =   array2table(Table_Data,'VariableNames',Table_Cols,'RowNames',Table_Rows);
    IrfAggTable.(SH).Diff_dom_ext ...
                =   IrfAggTable.(SH).dom-IrfAggTable.(SH).ext;
    IrfAggTable.(SH).Diff_N_H ...
                =   IrfAggTable.(SH).N-IrfAggTable.(SH).H;
    IrfAggTable.(SH).Diff_poor_rich ...
                =   IrfAggTable.(SH).poor-IrfAggTable.(SH).rich;
    % Pct
    IrfDecTable.(SH) ...
                =   IrfAggTable.(SH);
    for jj=1:length(IrfAggTable.(SH).Properties.VariableNames)
        vv          =   IrfAggTable.(SH).Properties.VariableNames{jj};
        IrfDecTable.(SH).(vv) ...
                    =   IrfDecTable.(SH).(vv)/sum(IrfDecTable.(SH).(vv));
    end
    % Print Tables
     for jj=1:length(IrfAggTable.(SH).Properties.VariableNames)
        vv          =   IrfAggTable.(SH).Properties.VariableNames{jj};
        IrfAggTable.(SH).(vv)   =   round(IrfAggTable.(SH).(vv),3);
        IrfDecTable.(SH).(vv)   =   round(IrfDecTable.(SH).(vv),3);
    end
    writetable(IrfAggTable.(SH),[Folder,'IrfAgg.xlsx'],'WriteRowNames',true,'Sheet',SH);
    writetable(IrfDecTable.(SH),[Folder,'IrfDec.xlsx'],'WriteRowNames',true,'Sheet',SH);
end


%% Visual Decomposition by Groups
N_col       =   1;
N_row       =   1;
fig_size    =   [N_col,N_row].*[0.75,1]*1/2; 
fig_gap     =   [0.05,0.08];
fig_VMargin =   [0.05,0.03];
fig_HMargin =   [0.10,0.02]; 

fig_SubPlot =   @(ii_vv)subtightplot(N_row,N_col,ii_vv,...
                                     fig_gap,fig_VMargin,fig_HMargin);
AxFontSize  =   8;
TickFmt     =   '%.2f';
XLabel      =   @(Label)xlabel(Label,'fontsize',8,'interpreter','latex');
YLabel      =   @(Label)ylabel(Label,'fontsize',8,'interpreter','latex');
LEG         =   @(PH,LabelList)legend(PH,LabelList,'fontsize',8,'interpreter','latex','location','northeast');


for ii=1:length(ShockList)
    SH = ShockList{ii};
    
    YNames = cell(length(IrfAggTable.(SH).Row),1);
    for jj=1:length(YNames); YNames{jj} = VarGroupLabel.(IrfAggTable.(SH).Row{jj}); end
    % Financial Integration
    TempData = [IrfAggTable.(SH).dom,IrfAggTable.(SH).ext,IrfAggTable.(SH).Diff_dom_ext];
    XNames = {'Non-integrated','Integrated','Difference'};

    Fig = FigureSetup([SH,', Dec-FI'],fig_size);
    ax = fig_SubPlot(1);
    [~,ax] = Plot_StackedBars(TempData,YNames,XNames,[],ax);
    
    print('-depsc','-r1000',Fig,[Folder,'byFI_',SH]);
    close(Fig);
    % Real Integration
    TempData = [IrfAggTable.(SH).N,IrfAggTable.(SH).H,IrfAggTable.(SH).Diff_N_H];
    XNames = {'Non-integrated','Integrated','Difference'};

    Fig = FigureSetup([SH,', Dec-RI'],fig_size);
    ax = fig_SubPlot(1);
    [~,ax] = Plot_StackedBars(TempData,YNames,XNames,[],ax);
    
    print('-depsc','-r1000',Fig,[Folder,'byRI_',SH]);
    close(Fig);
    % by Net Wealth
    TempData = [IrfAggTable.(SH).poor,IrfAggTable.(SH).rich,IrfAggTable.(SH).Diff_poor_rich];
    XNames = {'High Net Asset','Low Net Asset','Difference'};

    Fig = FigureSetup([SH,', Net Asset'],fig_size);
    ax = fig_SubPlot(1);
    [~,ax] = Plot_StackedBars(TempData,YNames,XNames,[],ax);
    
    print('-depsc','-r1000',Fig,[Folder,'byNA_',SH]);
    close(Fig);
end


